This PR adds in support for quantizing nn.Parameter to quantize_.

ModuleFqnToConfig has been renamed to FqnToConfig, which now accepts both module fqn and parameter fqns. ModuleFqnToConfig has been aliased to maintain BC.

bc-breaking changes

1. Passing filter_fn=None to quantize_ has changed semantics.

Previously, when we passed in filter_fn=None , we would just assign it to be _is_linear. So passing in None and _is_linear was functionally the same.

Now, None and _is_linear have different semantics. None will ignore filter_fn completely and just use the provided FqnToConfig to quantize the model.

before

# these are equivalent
model = nn.Linear(128, 128)

quantize(model, config, filter_fn=None) = quantize(mode, config, filter_fn=_is_linear)

assert isinstance(model.weight, TorchAOBaseTensor)

after

# these are equivalent
model = nn.Linear(128, 128)

quantize(model, config, filter_fn=None)
quantize(mode, config, filter_fn=_is_linear)

assert isinstance(model.weight, TorchAOBaseTensor)

This is needed because we have non-linear non-weight modules we want to quantize, and we need a module to both pass filter_fn and be specified in the FqnToConfig.

2. The default value of filter_fn for quantize_ has changed from None -> _is_linear.

To maintain the behavior of quantize(model, config), we have changed the default value of filter_fn from None to _is_linear explicitly.

3. Passing in filter_fn=None with _default in FqnToConfig now raises a ValueError.

Before this would default to _is_linear, but now that filter_fn=None has a different meaning _default is only supported when a filter_fn is specified.

4. Passing in non-default filter_fn with FqnToConfig will raise a ValueError

To encourage users to use avoid mixing filter_fn and FqnToConfig configuration, quantize_ will now throw a ValueError if filter_fn is not _is_linear or None and is a custom filter_fn.

API examples

For example, a toy nn.Linear model,

class MyModel(nn.Module()):
  linear1 = nn.Linear(128, 128)
  linear2 = nn.Linear(128, 128)

model = MyModel()

The keys to FqnToConfig can be one of the following (in order of precedence):

1. exact module FQN - We can quantize the weight of the first linear as follows

quant_config = FqnToConfig({
    "linear1": Float8DynamicActivationFloat8WeightConfig(
        granularity=PerRow(),
    ),
})

2. regex that matches module FQN (prepended by re:)

quant_config = FqnToConfig({
    "re:linear*": Float8DynamicActivationFloat8WeightConfig(
        granularity=PerRow(),
    ),
})

3. exact parameter FQN

quant_config = FqnToConfig({
    "linear1.weight": Float8DynamicActivationFloat8WeightConfig(
        granularity=PerRow(),
    ),
})

4. regex that matches parameter FQN (prepended by re:)

quant_config = FqnToConfig({
    "re:linear*.weight": Float8DynamicActivationFloat8WeightConfig(
        granularity=PerRow(),
    ),
})

5. _default, only when filter_fn is specified

quant_config = FqnToConfig({
    "_default": Float8DynamicActivationFloat8WeightConfig(
        granularity=PerRow(),
    ),
})

To enable support for parameter fqn for a paticular config, we need to add the parameter_name kwarg into the config signature, and update CUSTOM_PARAM_QUANTIZATION_SUPPOTED_CONFIGS. See the changes here for more details.

Float8DynamicActivationFloat8WeightConfig has been enabled by this PR, but other configs will throw an NotImplementedError.

Test Plan

1. unit tests for new config:

pytest test/quantization/test_quant_api.py::TestModuleOrParamFqnToConfig

2. regression test for ModuleFqnToConfig

pytest test/quantization/test_quant_api.py -k test_module_fqn_to_config

3. Make sure that we can load old HF checkpoints to maintain BC, run this

How do our configs translate for MoEs?

Currently, we define a bunch of configs that are for dense nn.Linear modules, how do these configs translate in the case of MoE inference?

Some background on MoE inference

There are two ways that forwards is implemented for MoE

• For loop of nn.Linear - In this case, we break down the 3d weight x activation matmul into a for loop of 2d weight x activation matmuls. This can be seen here.

In this case, I argue that the semantics of the configs do not change at all from the normal nn.Linear case, as we are just doing a bunch of normal 2d linear matmuls.

• bmm/grouped mm on the 3d weights / activations directly.

For this case, we'd need to add additional op support (bmm) for forwards. Depending on whether the subclass is an AQT subclass or non AQT subclass this will be added differently.

I plan to only support parameter quantization for non-AQT subclasses, my reasoning being that those are the most popular / important configs anyway (Float8Dynamic, Int4WeightOnly).

Below is a breakdown of what Configs map to AQT / non-AQT subclasses:

For these the majority of the semantics remain the same, the only semantics that really changes is PerRow granularity. and there's a very natural extension of PerRow to the 3d case (apply on the last dimension).

I took a look at the keys of the non-AQT configs below and what they would mean for MoEs.

Float8DynamicActivationFloat8WeightConfig

[('activation_dtype', <class 'torch.dtype'>),
 ('weight_dtype', <class 'torch.dtype'>),
 ('granularity',
  typing.Union[ForwardRef('PerTensor'), ForwardRef('PerRow'), typing.List[typing.Union[ForwardRef('PerTensor'), ForwardRef('PerRow')]], NoneType]),
 ('mm_config', typing.Optional[torchao.float8.inference.Float8MMConfig]),
 ('activation_value_lb', typing.Optional[float]),
 ('activation_value_ub', typing.Optional[float]),
 ('kernel_preference', <enum 'KernelPreference'>),
 ('set_inductor_config', <class 'bool'>),
 ('version', <class 'int'>)]

activation_dtype, weight_dtype, activation_value_lb, activation_value_ub all do not change meaning semantically.
granularity=PerTensor() does not change semantic meaning - we still use a single tensor to scale the entire weight tensor.
granularity=PerRow() does change meaning - we now calculate a scale for each row for the last dimension [-1] i.e for a weight of (E, N, K) we would expect PerRow to create scales of block size (1, 1, K).
mm_config kernel_preference and set_inductor_config stay the same as well.

Float8StaticActivationFloat8WeightConfig

[('scale', <class 'torch.Tensor'>),
 ('activation_dtype', <class 'torch.dtype'>),
 ('weight_dtype', <class 'torch.dtype'>),
 ('granularity',
  typing.Union[ForwardRef('PerTensor'), ForwardRef('PerRow'), typing.Tuple[typing.Union[ForwardRef('PerTensor'), ForwardRef('PerRow')], typing.Union[ForwardRef('PerTensor'), ForwardRef('PerRow')]], NoneType]),
 ('mm_config', typing.Optional[torchao.float8.inference.Float8MMConfig]),
 ('set_inductor_config', <class 'bool'>)]

scale should be passed in as a 3d tensor instead of a 2d tensor in the case of PerRow granularity

Float8DynamicActivationInt4WeightConfig

[('int4_packing_format', <enum 'Int4PackingFormat'>)]

int4_packing_format - Only "preshuffled" is supported and Int4PreshuffledTensor supports 3d weights.

Int4WeightOnlyConfig

[('group_size', <class 'int'>),
 ('layout',
  typing.Optional[torchao.dtypes.uintx.tensor_core_tiled_layout.TensorCoreTiledLayout]),
 ('use_hqq', <class 'bool'>),
 ('zero_point_domain',
  typing.Optional[torchao.quantization.quant_primitives.ZeroPointDomain]),
 ('set_inductor_config', <class 'bool'>),
 ('preserve_zero', typing.Optional[bool]),
 ('int4_packing_format', <enum 'Int4PackingFormat'>),
 ('int4_choose_qparams_algorithm', <enum 'Int4ChooseQParamsAlgorithm'>),
 ('version', <class 'int'>)]

group_size, int4_packing_format, int4_choose_qparams_algorithm, set_inductor_config are the only things that are set for v2 config,

I don't think these semantics of these change, although there are some packing formats that do not support 3d weights. It looks like (Int4PackingFormat.PLAIN_INT32, Int4PackingFormat.MARLIN_SPARSE).